from datetime import datetime

from database.settings import settings
from database.database import DBSessionManager
from models import TaskSite


class TaskSiteDAO:
    def __init__(self):
        # 获取全局配置的会话管理器配置
        self.session_config = settings.get_db_session_config(
            # 可以针对TaskSiteDAO进行特殊配置
            exclude_fields={"password", "api_key"}  # 排除敏感字段
        )

    def create_task_site(self, task_site: TaskSite):
        """创建任务站点"""
        with DBSessionManager(**self.session_config) as db:
            db.add(task_site)
            db.commit()
            # 注意：如果启用了序列化，refresh后需要重新查询
            if self.session_config.get('serialize_results', True):
                return db.get_by_id(TaskSite, task_site.id)
            else:
                db.refresh(task_site)
                return task_site

    def get_all_task_sites(self, type="new"):
        """获取所有任务站点（自动序列化）"""
        with DBSessionManager(**self.session_config) as db:
            return db.all(
                db.query(TaskSite)
                .filter(TaskSite.type == type, TaskSite.status == 1)
                .order_by(TaskSite.sort)
            )

    def get_task_by_id(self, id):
        """通过ID获取任务站点"""
        with DBSessionManager(**self.session_config) as db:
            return db.get_by_id(TaskSite, id)

    def update_site_execute_time(self, id):
        """
        更新最后采集的时间
        :param id: 站点ID
        :return: 更新的记录数
        """
        with DBSessionManager(**self.session_config) as db:
            # 使用更高效的方式更新（无需加载对象）
            result = db.query(TaskSite).filter(TaskSite.id == id).update(
                {"last_execute_time": datetime.now()},
                synchronize_session=False  # 禁用会话同步，提高性能
            )
            # 由于启用了自动提交，无需手动commit
            return result

    def batch_update_sites_status(self, site_ids, status):
        """
        批量更新站点状态（高效处理）
        :param site_ids: 站点ID列表
        :param status: 新状态值
        :return: 更新的记录数
        """
        with DBSessionManager(**self.session_config) as db:
            # 使用高效批量更新方法
            status_map = {site_id: status for site_id in site_ids}
            return db.batch_update_by_ids(
                TaskSite,
                "id",
                status_map,
                status_field="status"
            )

    def process_sites_in_batches(self, process_func):
        """
        流式处理所有站点（避免内存问题）
        :param process_func: 处理函数
        :return: 处理的记录数
        """
        with DBSessionManager(**self.session_config) as db:
            return db.process_in_batches(
                db.query(TaskSite),
                process_func,
                serialize=True
            )